"""
Redis状态存储后端 - 高性能状态存储

这个模块实现了基于Redis的状态存储：
- 高性能读写操作
- 数据持久化
- 分布式锁支持
- 过期时间管理
- 集群支持
"""

import asyncio
import json
import logging
import time
from datetime import datetime, timezone, timedelta
from typing import Any, Dict, List, Optional, Set, Union
import aioredis
from aioredis import Redis

from .result_aggregator import StateStorage, TaskRecord, TaskStatus

logger = logging.getLogger(__name__)


class RedisStateStorage(StateStorage):
    """Redis状态存储实现"""
    
    def __init__(self, 
                 redis_url: str = "redis://localhost:6379",
                 key_prefix: str = "gpu_auto_scaling",
                 default_ttl: int = 86400 * 7,  # 7天
                 connection_pool_size: int = 10):
        self.redis_url = redis_url
        self.key_prefix = key_prefix
        self.default_ttl = default_ttl
        self.connection_pool_size = connection_pool_size
        
        self.redis: Optional[Redis] = None
        self._connection_pool: Optional[aioredis.ConnectionPool] = None
        
        # 键名模板
        self.task_key_template = f"{key_prefix}:task:{{task_id}}"
        self.task_status_key_template = f"{key_prefix}:status:{{status}}"
        self.server_tasks_key_template = f"{key_prefix}:server:{{server_id}}:tasks"
        self.task_index_key = f"{key_prefix}:task_index"
        self.lock_key_template = f"{key_prefix}:lock:{{resource}}"
        
        self.logger = logging.getLogger(__name__)
        self.logger.info("Redis状态存储已初始化")
    
    async def connect(self) -> None:
        """连接到Redis"""
        try:
            # 创建连接池
            self._connection_pool = aioredis.ConnectionPool.from_url(
                self.redis_url,
                max_connections=self.connection_pool_size,
                retry_on_timeout=True
            )
            
            # 创建Redis客户端
            self.redis = aioredis.Redis(connection_pool=self._connection_pool)
            
            # 测试连接
            await self.redis.ping()
            
            self.logger.info("Redis连接成功")
            
        except Exception as e:
            self.logger.error(f"Redis连接失败: {e}")
            raise
    
    async def disconnect(self) -> None:
        """断开Redis连接"""
        if self.redis:
            await self.redis.close()
            self.redis = None
        
        if self._connection_pool:
            await self._connection_pool.disconnect()
            self._connection_pool = None
        
        self.logger.info("Redis连接已断开")
    
    async def get_task(self, task_id: str) -> Optional[TaskRecord]:
        """获取任务记录"""
        if not self.redis:
            await self.connect()
        
        try:
            task_key = self.task_key_template.format(task_id=task_id)
            task_data = await self.redis.get(task_key)
            
            if not task_data:
                return None
            
            # 解析任务数据
            task_dict = json.loads(task_data)
            return self._dict_to_task_record(task_dict)
            
        except Exception as e:
            self.logger.error(f"获取任务记录失败: {task_id}, 错误: {e}")
            return None
    
    async def update_task(self, task_record: TaskRecord) -> None:
        """更新任务记录"""
        if not self.redis:
            await self.connect()
        
        try:
            task_key = self.task_key_template.format(task_id=task_record.task_id)
            
            # 获取旧的任务记录（如果存在）
            old_task_data = await self.redis.get(task_key)
            old_task_record = None
            if old_task_data:
                old_task_dict = json.loads(old_task_data)
                old_task_record = self._dict_to_task_record(old_task_dict)
            
            # 转换为字典格式
            task_dict = self._task_record_to_dict(task_record)
            task_json = json.dumps(task_dict)
            
            # 使用管道进行批量操作
            pipe = self.redis.pipeline()
            
            # 保存任务记录
            pipe.setex(task_key, self.default_ttl, task_json)
            
            # 更新状态索引
            if old_task_record and old_task_record.status != task_record.status:
                # 从旧状态集合中移除
                old_status_key = self.task_status_key_template.format(status=old_task_record.status.value)
                pipe.srem(old_status_key, task_record.task_id)
            
            # 添加到新状态集合
            status_key = self.task_status_key_template.format(status=task_record.status.value)
            pipe.sadd(status_key, task_record.task_id)
            
            # 更新服务器任务索引
            if task_record.assigned_server:
                server_tasks_key = self.server_tasks_key_template.format(
                    server_id=task_record.assigned_server
                )
                pipe.sadd(server_tasks_key, task_record.task_id)
                pipe.expire(server_tasks_key, self.default_ttl)
            
            # 更新任务索引
            pipe.zadd(self.task_index_key, {task_record.task_id: task_record.submit_time.timestamp()})
            pipe.expire(self.task_index_key, self.default_ttl)
            
            # 执行管道操作
            await pipe.execute()
            
            self.logger.debug(f"已更新任务记录: {task_record.task_id}")
            
        except Exception as e:
            self.logger.error(f"更新任务记录失败: {task_record.task_id}, 错误: {e}")
            raise
    
    async def list_tasks(self, 
                        status: Optional[TaskStatus] = None,
                        server_id: Optional[str] = None,
                        limit: int = 100) -> List[TaskRecord]:
        """列出任务记录"""
        if not self.redis:
            await self.connect()
        
        try:
            task_ids = set()
            
            if status:
                # 从状态索引获取任务ID
                status_key = self.task_status_key_template.format(status=status.value)
                status_task_ids = await self.redis.smembers(status_key)
                task_ids.update(task_id.decode() for task_id in status_task_ids)
            
            if server_id:
                # 从服务器索引获取任务ID
                server_tasks_key = self.server_tasks_key_template.format(server_id=server_id)
                server_task_ids = await self.redis.smembers(server_tasks_key)
                server_task_ids = set(task_id.decode() for task_id in server_task_ids)
                
                if task_ids:
                    task_ids = task_ids.intersection(server_task_ids)
                else:
                    task_ids = server_task_ids
            
            if not task_ids and not status and not server_id:
                # 如果没有过滤条件，从任务索引获取最新的任务
                task_index_data = await self.redis.zrevrange(
                    self.task_index_key, 0, limit - 1, withscores=False
                )
                task_ids = set(task_id.decode() for task_id in task_index_data)
            
            # 限制结果数量
            task_ids = list(task_ids)[:limit]
            
            # 批量获取任务记录
            if not task_ids:
                return []
            
            pipe = self.redis.pipeline()
            for task_id in task_ids:
                task_key = self.task_key_template.format(task_id=task_id)
                pipe.get(task_key)
            
            task_data_list = await pipe.execute()
            
            # 解析任务记录
            tasks = []
            for task_data in task_data_list:
                if task_data:
                    task_dict = json.loads(task_data)
                    task_record = self._dict_to_task_record(task_dict)
                    tasks.append(task_record)
            
            return tasks
            
        except Exception as e:
            self.logger.error(f"列出任务记录失败: {e}")
            return []
    
    async def delete_task(self, task_id: str) -> bool:
        """删除任务记录"""
        if not self.redis:
            await self.connect()
        
        try:
            # 获取任务记录以获取索引信息
            task_record = await self.get_task(task_id)
            if not task_record:
                return False
            
            # 使用管道进行批量删除
            pipe = self.redis.pipeline()
            
            # 删除任务记录
            task_key = self.task_key_template.format(task_id=task_id)
            pipe.delete(task_key)
            
            # 从状态索引中移除
            status_key = self.task_status_key_template.format(status=task_record.status.value)
            pipe.srem(status_key, task_id)
            
            # 从服务器索引中移除
            if task_record.assigned_server:
                server_tasks_key = self.server_tasks_key_template.format(
                    server_id=task_record.assigned_server
                )
                pipe.srem(server_tasks_key, task_id)
            
            # 从任务索引中移除
            pipe.zrem(self.task_index_key, task_id)
            
            # 执行删除操作
            await pipe.execute()
            
            self.logger.info(f"已删除任务记录: {task_id}")
            return True
            
        except Exception as e:
            self.logger.error(f"删除任务记录失败: {task_id}, 错误: {e}")
            return False
    
    async def get_task_count(self, status: Optional[TaskStatus] = None) -> int:
        """获取任务数量"""
        if not self.redis:
            await self.connect()
        
        try:
            if status:
                status_key = self.task_status_key_template.format(status=status.value)
                return await self.redis.scard(status_key)
            else:
                return await self.redis.zcard(self.task_index_key)
                
        except Exception as e:
            self.logger.error(f"获取任务数量失败: {e}")
            return 0
    
    async def get_server_task_count(self, server_id: str) -> int:
        """获取服务器任务数量"""
        if not self.redis:
            await self.connect()
        
        try:
            server_tasks_key = self.server_tasks_key_template.format(server_id=server_id)
            return await self.redis.scard(server_tasks_key)
            
        except Exception as e:
            self.logger.error(f"获取服务器任务数量失败: {server_id}, 错误: {e}")
            return 0
    
    async def cleanup_expired_tasks(self, max_age_days: int = 30) -> int:
        """清理过期任务"""
        if not self.redis:
            await self.connect()
        
        try:
            cutoff_timestamp = (datetime.now(timezone.utc) - timedelta(days=max_age_days)).timestamp()
            
            # 获取过期的任务ID
            expired_task_ids = await self.redis.zrangebyscore(
                self.task_index_key, 0, cutoff_timestamp
            )
            
            if not expired_task_ids:
                return 0
            
            # 批量删除过期任务
            deleted_count = 0
            for task_id_bytes in expired_task_ids:
                task_id = task_id_bytes.decode()
                if await self.delete_task(task_id):
                    deleted_count += 1
            
            self.logger.info(f"已清理 {deleted_count} 个过期任务")
            return deleted_count
            
        except Exception as e:
            self.logger.error(f"清理过期任务失败: {e}")
            return 0
    
    async def acquire_lock(self, resource: str, timeout: int = 30) -> bool:
        """获取分布式锁"""
        if not self.redis:
            await self.connect()
        
        try:
            lock_key = self.lock_key_template.format(resource=resource)
            lock_value = f"{time.time()}_{id(self)}"
            
            # 尝试获取锁
            result = await self.redis.set(
                lock_key, 
                lock_value, 
                nx=True,  # 只在键不存在时设置
                ex=timeout  # 设置过期时间
            )
            
            return result is not None
            
        except Exception as e:
            self.logger.error(f"获取锁失败: {resource}, 错误: {e}")
            return False
    
    async def release_lock(self, resource: str) -> bool:
        """释放分布式锁"""
        if not self.redis:
            await self.connect()
        
        try:
            lock_key = self.lock_key_template.format(resource=resource)
            result = await self.redis.delete(lock_key)
            return result > 0
            
        except Exception as e:
            self.logger.error(f"释放锁失败: {resource}, 错误: {e}")
            return False
    
    async def get_statistics(self) -> Dict[str, Any]:
        """获取存储统计信息"""
        if not self.redis:
            await self.connect()
        
        try:
            stats = {
                "total_tasks": await self.get_task_count(),
                "status_counts": {},
                "server_counts": {},
                "redis_info": {}
            }
            
            # 获取各状态的任务数量
            for status in TaskStatus:
                count = await self.get_task_count(status)
                stats["status_counts"][status.value] = count
            
            # 获取Redis信息
            redis_info = await self.redis.info()
            stats["redis_info"] = {
                "used_memory": redis_info.get("used_memory_human"),
                "connected_clients": redis_info.get("connected_clients"),
                "total_commands_processed": redis_info.get("total_commands_processed"),
                "keyspace_hits": redis_info.get("keyspace_hits"),
                "keyspace_misses": redis_info.get("keyspace_misses")
            }
            
            return stats
            
        except Exception as e:
            self.logger.error(f"获取统计信息失败: {e}")
            return {}
    
    def _task_record_to_dict(self, task_record: TaskRecord) -> Dict[str, Any]:
        """将任务记录转换为字典"""
        return {
            "task_id": task_record.task_id,
            "client_ip": task_record.client_ip,
            "submit_time": task_record.submit_time.isoformat(),
            "assigned_server": task_record.assigned_server,
            "server_endpoint": task_record.server_endpoint,
            "status": task_record.status.value,
            "result_url": task_record.result_url,
            "processing_duration": task_record.processing_duration,
            "error_message": task_record.error_message,
            "correlation_id": task_record.correlation_id,
            "metadata": task_record.metadata or {}
        }
    
    def _dict_to_task_record(self, task_dict: Dict[str, Any]) -> TaskRecord:
        """将字典转换为任务记录"""
        return TaskRecord(
            task_id=task_dict["task_id"],
            client_ip=task_dict["client_ip"],
            submit_time=datetime.fromisoformat(task_dict["submit_time"]),
            assigned_server=task_dict.get("assigned_server"),
            server_endpoint=task_dict.get("server_endpoint"),
            status=TaskStatus(task_dict["status"]),
            result_url=task_dict.get("result_url"),
            processing_duration=task_dict.get("processing_duration"),
            error_message=task_dict.get("error_message"),
            correlation_id=task_dict.get("correlation_id"),
            metadata=task_dict.get("metadata", {})
        )
